package rediss

import (
	"context"
	"github.com/go-redis/redis/extra/rediscmd/v8"
	"github.com/go-redis/redis/v8"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
)

type OpentracingHook struct {
}

func NewTracingHook() *OpentracingHook {
	return new(OpentracingHook)
}

func (o OpentracingHook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (context.Context, error) {
	span, newCtx := opentracing.StartSpanFromContext(ctx, cmd.FullName())
	ext.DBType.Set(span, "redis")
	ext.DBStatement.Set(span, rediscmd.CmdString(cmd))
	return newCtx, nil
}

func (o OpentracingHook) AfterProcess(ctx context.Context, cmd redis.Cmder) error {
	span := opentracing.SpanFromContext(ctx)
	if err := cmd.Err(); err != nil {
		recordError(ctx, span, err)
	}
	span.Finish()
	return nil
}

func (o OpentracingHook) BeforeProcessPipeline(ctx context.Context, cmds []redis.Cmder) (context.Context, error) {
	summary, cmdsString := rediscmd.CmdsString(cmds)

	span, ctx := opentracing.StartSpanFromContext(ctx, "pipeline "+summary)
	ext.DBType.Set(span, "redis")
	ext.DBStatement.Set(span, cmdsString)
	return ctx, nil
}

func (o OpentracingHook) AfterProcessPipeline(ctx context.Context, cmds []redis.Cmder) error {
	span := opentracing.SpanFromContext(ctx)
	if err := cmds[0].Err(); err != nil {
		recordError(ctx, span, err)
	}
	span.Finish()
	return nil
}

func recordError(ctx context.Context, span opentracing.Span, err error) {
	if err != redis.Nil {
		ext.Error.Set(span, true)
	}
}
